import os
from langchain_core.messages import HumanMessage
from dotenv import load_dotenv
from langgraph.checkpoint.memory import MemorySaver
from langgraph.graph import END, StateGraph, MessagesState
from langgraph.prebuilt import ToolNode
from langchain_core.tools import tool
from langchain_deepseek import ChatDeepSeek
from IPython.display import Image, display

load_dotenv("./config/config.env")

# 从环境变量中获取配置值
langsmith_api_key = os.getenv("LANGSMITH_API_KEY")

# 设置 LangSmith 环境变量
os.environ["LANGCHAIN_TRACING_V2"] = "true"
os.environ["LANGCHAIN_API_KEY"] = langsmith_api_key if langsmith_api_key else "your-langsmith-api-key"
os.environ["LANGCHAIN_PROJECT"] = "weather-agent-workflow"

model = "deepseek-chat"
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")

llm = ChatDeepSeek(model=model, temperature=0, max_retries=2, api_key=OPENAI_API_KEY)

# 1. 定义工具
@tool
def search(query: str):
    """模拟一个搜索工具"""
    if "上海" in query:
        return "现在30度，有雾."
    return "现在是35度，阳光明媚。"

tools = [search]
llm_with_tools = llm.bind_tools(tools)
tool_node = ToolNode(tools)

# 3. 定义状态和节点函数
class State(MessagesState):
    pass

def should_continue(state: MessagesState):
    messages = state['messages']
    last_message = messages[-1]
    if last_message.tool_calls:
        return "tools"
    return END

def call_model(state: MessagesState):
    messages = state['messages']
    response = llm_with_tools.invoke(messages)
    return {"messages": [response]}

# 4. 构建图
workflow = StateGraph(State)
workflow.add_node("agent", call_model)
workflow.add_node("tools", tool_node)
workflow.set_entry_point("agent")

workflow.add_conditional_edges("agent", should_continue)
workflow.add_edge("tools", "agent")

# 5. 编译图
checkpointer = MemorySaver()
app = workflow.compile(checkpointer=checkpointer)

# 6. 执行图
final_state = app.invoke(
    {"messages": [HumanMessage(content="上海的天气怎么样?")]},
    config={
        "configurable": {"thread_id": 42},
        "metadata": {"langsmith_project": "weather-agent-workflow"}
    }
)
print(final_state["messages"][-1].content)

# 7. 可视化图结构 - 简单版本
# 最简单的备选方案
try:
    # 保存为PNG文件
    png_data = app.get_graph().draw_mermaid_png()
    with open("langgraph_workflow.png", "wb") as f:
        f.write(png_data)
    print("图表已保存为 langgraph_workflow.png")
except Exception as e:
    print(f"保存图表失败: {e}")
    # 打印 mermaid 代码作为备选
    try:
        print("Mermaid代码:")
        print(app.get_graph().draw_mermaid())
    except:
        pass